// Copyright Epic Games, Inc. All Rights Reserved.

#pragma once

#include <zencore/compactbinarypackage.h>
#include <zencore/iobuffer.h>
#include <zencore/iohash.h>

#include <functional>
#include <gsl/gsl-lite.hpp>

namespace zen {

class IoBuffer;
class CbPackage;
class CompositeBuffer;

/**	   _____ _     _____           _
	  / ____| |   |  __ \         | |
	 | |    | |__ | |__) |_ _  ___| | ____ _  __ _  ___
	 | |    | '_ \|  ___/ _` |/ __| |/ / _` |/ _` |/ _ \
	 | |____| |_) | |  | (_| | (__|   < (_| | (_| |  __/
	  \_____|_.__/|_|   \__,_|\___|_|\_\__,_|\__, |\___|
											  __/ |
											 |___/

	Structures and code related to handling CbPackage transactions

	CbPackage instances are marshaled across the wire using a distinct message
	format. We don't use the CbPackage serialization format provided by the
	CbPackage implementation itself since that does not provide much flexibility
	in how the attachment payloads are transmitted. The scheme below separates
	metadata cleanly from payloads and this enables us to more efficiently
	transmit them either via sendfile/TransmitFile like mechanisms, or by
	reference/memory mapping in the local case.
 */

struct CbPackageHeader
{
	uint32_t HeaderMagic;
	uint32_t AttachmentCount;  // TODO: should add ability to opt out of implicit root document?
	uint32_t Reserved1;
	uint32_t Reserved2;
};

static_assert(sizeof(CbPackageHeader) == 16);

enum : uint32_t
{
	kCbPkgMagic = 0xaa77aacc
};

struct CbAttachmentEntry
{
	uint64_t PayloadSize;	  // Size of the associated payload data in the message
	uint32_t Flags;			  // See flags below
	IoHash	 AttachmentHash;  // Content Id for the attachment

	enum
	{
		kIsCompressed = (1u << 0),	// Is marshaled using compressed buffer storage format
		kIsObject	  = (1u << 1),	// Is compact binary object
		kIsError	  = (1u << 2),	// Is error (compact binary formatted) object
		kIsLocalRef	  = (1u << 3),	// Is "local reference"
	};
};

struct CbAttachmentReferenceHeader
{
	uint64_t PayloadByteOffset	= 0;
	uint64_t PayloadByteSize	= ~0u;
	uint16_t AbsolutePathLength = 0;

	// This header will be followed by UTF8 encoded absolute path to backing file
};

static_assert(sizeof(CbAttachmentEntry) == 32);

enum class FormatFlags
{
	kDefault					= 0,
	kAllowLocalReferences		= (1u << 0),
	kDenyPartialLocalReferences = (1u << 1)
};

gsl_DEFINE_ENUM_BITMASK_OPERATORS(FormatFlags);

enum class RpcAcceptOptions : uint16_t
{
	kNone						 = 0,
	kAllowLocalReferences		 = (1u << 0),
	kAllowPartialLocalReferences = (1u << 1),
	kAllowPartialCacheChunks	 = (1u << 2)
};

gsl_DEFINE_ENUM_BITMASK_OPERATORS(RpcAcceptOptions);

std::vector<IoBuffer> FormatPackageMessage(const CbPackage& Data, FormatFlags Flags, void* TargetProcessHandle = nullptr);
CompositeBuffer		  FormatPackageMessageBuffer(const CbPackage& Data, FormatFlags Flags, void* TargetProcessHandle = nullptr);
CbPackage			  ParsePackageMessage(
				IoBuffer												  Payload,
				std::function<IoBuffer(const IoHash& Cid, uint64_t Size)> CreateBuffer = [](const IoHash&, uint64_t Size) -> IoBuffer {
		return IoBuffer{Size};
				});
bool IsPackageMessage(IoBuffer Payload);

bool ParsePackageMessageWithLegacyFallback(const IoBuffer& Response, CbPackage& OutPackage);

std::vector<IoBuffer> FormatPackageMessage(const CbPackage& Data, void* TargetProcessHandle = nullptr);
CompositeBuffer		  FormatPackageMessageBuffer(const CbPackage& Data, void* TargetProcessHandle = nullptr);

/** Streaming reader for compact binary packages

	The goal is to ultimately support zero-copy I/O, but for now there'll be some
	copying involved on some platforms at least.

	This approach to deserializing CbPackage data is more efficient than
	`ParsePackageMessage` since it does not require the entire message to
	be resident in a memory buffer

  */
class CbPackageReader
{
public:
	CbPackageReader();
	~CbPackageReader();

	void SetPayloadBufferCreator(std::function<IoBuffer(const IoHash& Cid, uint64_t Size)> CreateBuffer);

	/** Process compact binary package data stream

		The data stream must be in the serialization format produced by FormatPackageMessage

		\return How many bytes must be fed to this function in the next call
	 */
	uint64_t ProcessPackageHeaderData(const void* Data, uint64_t DataBytes);

	void							 Finalize();
	const std::vector<CbAttachment>& GetAttachments() { return m_Attachments; }
	CbObject						 GetRootObject() { return m_RootObject; }
	std::span<IoBuffer>				 GetPayloadBuffers() { return m_PayloadBuffers; }

private:
	enum class State
	{
		kInitialState,
		kReadingHeader,
		kReadingAttachmentEntries,
		kReadingBuffers
	} m_CurrentState = State::kInitialState;

	std::function<IoBuffer(const IoHash& Cid, uint64_t Size)> m_CreateBuffer;
	std::vector<IoBuffer>									  m_PayloadBuffers;
	std::vector<CbAttachmentEntry>							  m_AttachmentEntries;
	std::vector<CbAttachment>								  m_Attachments;
	CbObject												  m_RootObject;
	CbPackageHeader											  m_PackageHeader;

	IoBuffer MarshalLocalChunkReference(IoBuffer AttachmentBuffer);
};

void forcelink_packageformat();

}  // namespace zen
